#include "oqs_prov.h"
#include "test_common.h"
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <openssl/core_dispatch.h>
#include <openssl/core_names.h>
#include <openssl/provider.h>
#include <openssl/ssl.h>
#include <openssl/pem.h>
#include <openssl/x509.h>
#include <openssl/x509v3.h>
#include <openssl/trace.h>
#include <openssl/params.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/param_build.h>
#include <openssl/bio.h> 
#include <openssl/asn1.h>
#include <openssl/ec.h>
#include <openssl/rsa.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <time.h> 
const char *kHybridSignatureAlgorithms[] = {
    "p256_dilithium2", "rsa3072_dilithium2", "p384_dilithium3", "p521_dilithium5",
    "p256_mldsa44", "rsa3072_mldsa44", "p384_mldsa65", "p521_mldsa87",
    "p256_falcon512", "rsa3072_falcon512", "p256_falconpadded512", "rsa3072_falconpadded512",
    "p521_falcon1024", "p521_falconpadded1024", "p256_sphincssha2128fsimple", "rsa3072_sphincssha2128fsimple",
    "p256_sphincssha2128ssimple", "rsa3072_sphincssha2128ssimple", "p384_sphincssha2192fsimple",
    "p256_sphincsshake128fsimple", "rsa3072_sphincsshake128fsimple", NULL
};

struct SigPair {
    unsigned char *sigma1;
    size_t sigma1_len;
    unsigned char *sigma2;
    size_t sigma2_len;
} signature;

static int nid = NID_undef; 

void initialize_openssl_globals() {  
    nid = OBJ_create("1.2.3.4.1.2.3.4", "sig", "Mix_Sig");  
    if (nid == NID_undef) {  
    } 
}  

static int is_signature_algorithm_hybrid(const char **list, const char *s) {
    for (; *list != NULL && strcmp(*list, s) != 0; ++list);
    return *list != NULL;
}

static OSSL_LIB_CTX *init_openssl(void) {
    OSSL_LIB_CTX *ctx;

    if (!(ctx = OSSL_LIB_CTX_new())) {
        fprintf(stderr, cRED "failed to initialize a new `OSSL_LIB_CTX` " cNORM "\n");
    }
    return ctx;
}

char *get_oqs_fromtls(const char *tlsname){
    char *oqsname;
    char *del = "_";
    char *position = strchr(tlsname, *del);
    
    if(position == NULL)
    return NULL;
    
    oqsname = position+1;
    
    return oqsname;
}

char *get_cls_fromtls(const char *tlsname) {
    char *clsname;
    char *del = "_";
    char *position = strchr(tlsname, *del);

    if (position == NULL) {
           return NULL;
    }

    size_t len = position - tlsname;

    clsname = (char *)malloc(len + 1);  // +1 为了 '\0'

    if (clsname == NULL) {
        return NULL;
    }
    
    strncpy(clsname, tlsname, len);
    clsname[len] = '\0';

    return clsname;
}

X509_REQ *generate_cert_req(X509 *cert,EVP_PKEY *pkey){
    X509_REQ *cert_req = X509_REQ_new();
    X509_NAME *name = NULL;
          if(!cert_req){
       fprintf(stderr, "X509_REQ_new error.\n");
        X509_REQ_free(cert_req);
        X509_NAME_free(name);
                     return NULL;
    }
   if(!X509_REQ_set_pubkey(cert_req, pkey)||!(name = X509_REQ_get_subject_name(cert_req))||!X509_NAME_add_entry_by_txt(name, "C",  MBSTRING_ASC, (unsigned char *)"CN", -1, -1, 0)||!X509_NAME_add_entry_by_txt(name, "O",  MBSTRING_ASC, (unsigned char *)"icatest.org", -1, -1, 0)||!X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char *)"ICA", -1, -1, 0)||!    X509_REQ_sign(cert_req, pkey, EVP_sha256())){
          fprintf(stderr, "X509_REQ_gen error.\n");
          X509_REQ_free(cert_req);
          X509_NAME_free(name);
                          return NULL;
    }
    
        return cert_req;  
} 

 EVP_PKEY *read_pkey(const char *path){
     FILE *fp = fopen(path, "r");
          if (!fp) {
         perror("Error opening file");
                      return NULL;
     }
    EVP_PKEY *pkey = PEM_read_PrivateKey(fp, NULL, NULL, NULL);  
           if(!pkey){
         ERR_print_errors_fp(stderr);
         EVP_PKEY_free(pkey);
         fclose(fp);
                       return NULL;
    }
    
    fclose(fp);  
          return pkey;
 }
 
 X509 *read_certificate(const char *path) {
    FILE *fp = fopen(path, "r");
          if (!fp) {
         perror("Error opening file");
                       return NULL;
     }

     X509 *cert = PEM_read_X509(fp, NULL, NULL, NULL);
             if (!cert) {
         ERR_print_errors_fp(stderr);
         X509_free(cert);
         fclose(fp);
                        return NULL;
     }
     fclose(fp);
             return cert;
}

static int verify_sigma2(X509 *cert, EVP_PKEY *pkey){
          int  ret = 0;
          const ASN1_BIT_STRING *sig;
          const X509_ALGOR *sigalg = NULL;  
    X509_get0_signature(&sig, &sigalg, cert);
          size_t siglen = 0;
    siglen = sig->length;
    //fprintf(stderr, "Sig length is %d.\n", siglen);
    signature.sigma1_len = siglen;
    signature.sigma1 = (unsigned char *)malloc(siglen);
         if (signature.sigma1== NULL) {
        fprintf(stderr, "Memory allocation failed\n");
                     return ret ;
    }
    memcpy(signature.sigma1,sig->data, siglen); 
    
          unsigned char*tbs = NULL;
          size_t tbs_len = i2d_re_X509_tbs(cert, &tbs);
    //fprintf(stderr, "TBS length is %d.\n", tbs_len);
          
    EVP_MD_CTX *mdctx =  EVP_MD_CTX_new();
    EVP_DigestVerifyInit(mdctx, NULL, EVP_sha256(),NULL, pkey);
    EVP_DigestVerifyUpdate(mdctx, tbs, tbs_len);
    ret = EVP_DigestVerifyFinal(mdctx, signature.sigma1, signature.sigma1_len);    
          if (ret == 1) {  
       printf("Sigma2 verification success!\n");  
    } else {  
        printf("Sigma2 verification failed!\n");  
    } 
    EVP_MD_CTX_free(mdctx); 
    ret = 1;
          return ret;

}

static int verify_sigma1(X509 *cert, EVP_PKEY *pkey, X509_REQ *cert_req){
//  
          int ret = 0;
          int index = X509_get_ext_by_NID(cert, NID_subject_alt_name , -1);
          if(index == -1){
        fprintf(stderr, "Failed to find nid from ext.\n");
                      return ret;         
    }  
    
    X509_EXTENSION *ext = X509_get_ext(cert, index);
    ASN1_OCTET_STRING *data2 = X509_EXTENSION_get_data(ext);
          size_t Sigma1_len = data2->length;
           unsigned char*Sigma1 = (unsigned char*)malloc(Sigma1_len);
    memcpy(Sigma1, data2->data, Sigma1_len);
    
    X509 *x509 = X509_new();
    X509_NAME *name = X509_NAME_new();
    
   if(!X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char *)"CA", -1, -1, 0)||
  !X509_NAME_add_entry_by_txt(name, "O",  MBSTRING_ASC, (unsigned char *)"catest.org", -1, -1, 0)||!X509_NAME_add_entry_by_txt(name, "C",  MBSTRING_ASC, (unsigned char *)"CN", -1, -1, 0)){
        fprintf(stderr, "X509 NAME gen error.\n");
                     return ret;
    }          
    ASN1_INTEGER_set(X509_get_serialNumber(x509),1);
    X509_gmtime_adj(X509_getm_notBefore(x509), 0);
    X509_gmtime_adj(X509_getm_notAfter(x509), 31536000L);
          
    X509_set_subject_name(x509, X509_REQ_get_subject_name(cert_req));
    X509_set_issuer_name(cert, name);    
    X509_set_pubkey(x509, X509_REQ_get_pubkey(cert_req)); 
    
          if(!X509_sign(x509, pkey,EVP_sha256())){
        fprintf(stderr, "Sigma1_cert_gen error.\n");
        X509_free(x509);
                     return ret;
    }
 //   fprintf(stderr, "Sigma1_cert_gen success.\n");
    //
          const ASN1_BIT_STRING *sig;
          const X509_ALGOR *sigalg = NULL;  
    X509_get0_signature(&sig, &sigalg, x509);
    
          size_t siglen = 0;
    siglen = sig->length;
   // fprintf(stderr, "Sig length is %d.\n", siglen);
           unsigned char *sigma1 = NULL;
           sigma1 = (unsigned char *)malloc(siglen);
         if (sigma1== NULL) {
        fprintf(stderr, "Memory allocation failed\n");
                  return 0;
    }
    memcpy(sigma1,sig->data, siglen);    
          unsigned char*tbs = NULL;
          size_t tbs_len = i2d_re_X509_tbs(x509, &tbs);
  //  fprintf(stderr, "TBS length is %d.\n", tbs_len);
    
             
    EVP_MD_CTX *mdctx =  EVP_MD_CTX_new();
    EVP_DigestVerifyInit(mdctx, NULL, EVP_sha256(),NULL, pkey);
    EVP_DigestVerifyUpdate(mdctx, tbs, tbs_len);
    ret = EVP_DigestVerifyFinal(mdctx, sigma1, siglen);
    
    
          if (ret == 1) {  
        printf("Sigma1 verification success!\n");  
    } else {  
        printf("Sigma1 verification failed!\n");  
    }
          
          return ret;
}

int test_algorithm(OSSL_LIB_CTX *libctx,const char *signames,char * key_addr){

          int ret = 1;
    EVP_PKEY *class_pkey = NULL;
    EVP_PKEY *pq_pkey = NULL;
    X509 *cert = X509_new();
          char classprivkeypath[300];
          char oqsprivkeypath[300];
          char certpath[300];
          char keypath[300];
          const char *sep = "/";   

    sprintf(classprivkeypath, "%s%s%s%s", key_addr, sep, signames, "_cls.key");
    sprintf(oqsprivkeypath, "%s%s%s%s", key_addr, sep, signames, "_oqs.key");
    sprintf(certpath, "%s%s%s%s", key_addr, sep, signames, ".crt");
    
    cert = read_certificate(certpath);
    class_pkey = read_pkey(classprivkeypath);
    pq_pkey = read_pkey(oqsprivkeypath);
      
    signature.sigma1 = NULL;
    signature.sigma2 = NULL;
    signature.sigma1_len = 0;
    signature.sigma2_len = 0;
    
    X509_REQ *cert_req = X509_REQ_new();
    cert_req = generate_cert_req(cert, pq_pkey); 
    clock_t start_time = clock();      
    verify_sigma2(cert, pq_pkey);    
    
    verify_sigma1(cert, class_pkey , cert_req);

    clock_t end_time = clock(); 
    
    clock_t total_time = end_time - start_time;  
          double microseconds = ((double)total_time) * (1000000.0 / CLOCKS_PER_SEC);          
    fprintf(stderr, "程序运行时间: %.0f 微秒\n", microseconds);
   
    ret = 0;
end:
    EVP_PKEY_free(class_pkey);
    EVP_PKEY_free(pq_pkey);
        return ret;
}


int main(int argc, char *argv[]) {
  OSSL_LIB_CTX *libctx;
  OSSL_PROVIDER *oqs_provider;
    char key_addr[100] ;
    const char *modulename;
    const char *configfile;
    const OSSL_ALGORITHM *sigalgs;
    int query_nocache;  
    int ret = EXIT_FAILURE;
          
    if (!(libctx = init_openssl())) {
        goto end;
    }

    T(argc == 3);
    modulename = argv[1];
    configfile = argv[2];
    
    if(!(oqs_provider = OSSL_PROVIDER_load(libctx, modulename))){
        fputs(cRED "  `oqs_provider` is NULL " cNORM "\n", stderr);
        goto end;    
    }
   initialize_openssl_globals();   
   sigalgs = OSSL_PROVIDER_query_operation(oqs_provider, OSSL_OP_SIGNATURE,
                                            &query_nocache);
    
    fprintf(stdout, "Enter the path to load the files:");
        if (fgets(key_addr, sizeof(key_addr), stdin) == NULL) {
        fprintf(stderr, cRED " Failed to read input." cNORM "\n");
                  goto unload_def_provider;
    }
    key_addr[strcspn(key_addr, "\n")] = 0;
    //printf("%s\n", key_addr);  
    
    if (sigalgs) {
       for (; sigalgs->algorithm_names != NULL; sigalgs++) {
                  if (!is_signature_algorithm_hybrid(kHybridSignatureAlgorithms, sigalgs->algorithm_names)) {
           //fprintf(stderr, cRED "%s is not hybrid." cNORM "\n",sigalgs->algorithm_names);
                                continue;
          }
                            if (test_algorithm(libctx,sigalgs->algorithm_names, key_addr)) {
           fprintf(stderr, cRED " Failed for %s " cNORM "\n", sigalgs->algorithm_names);
       } else {
           fprintf(stderr, cGREEN "%s succeeded" cNORM "\n", sigalgs->algorithm_names);       
      }
    }
    }
    ret = 0;
unload_def_provider:
    OSSL_PROVIDER_unload(oqs_provider);   
end:
    OSSL_LIB_CTX_free(libctx);
    return ret;
  
}
